#include "kernel_operator.h"
using namespace AscendC;

class NLLLossSca {
 public:
  __aicore__ inline NLLLossSca() {}
  __aicore__ inline void Init(GM_ADDR x, GM_ADDR target, GM_ADDR weight,
                              GM_ADDR y, uint32_t reduction,
                              int32_t ignoreIndex, uint32_t sizeN,
                              uint32_t sizeC, uint32_t sizeM) {
    // printf("------------NLLLossSca------------\n");
    this->reduction = reduction;
    this->ignoreIndex = ignoreIndex;
    this->sizeN = sizeN;
    this->sizeC = sizeC;
    this->sizeM = sizeM;

    xGm.SetGlobalBuffer((__gm__ DTYPE_X*)x,
                        sizeN * sizeC * sizeM * sizeof(DTYPE_X));
    targetGm.SetGlobalBuffer((__gm__ DTYPE_TARGET*)target,
                             sizeN * sizeM * sizeof(DTYPE_TARGET));
    weightGm.SetGlobalBuffer((__gm__ DTYPE_WEIGHT*)weight,
                             sizeC * sizeof(DTYPE_WEIGHT));
    yGm.SetGlobalBuffer((__gm__ DTYPE_Y*)y, sizeof(DTYPE_Y));
    // printf("------------init down------------\n");
  }

  __aicore__ inline void process() {
    // printf("------------NLLLossSca process------------\n");
    float sumWeight = 0;
    float sumRes = 0;
    float tmpWeight, tmpSelect;
    uint32_t tmpTarget;
    uint32_t strideN = this->sizeC * this->sizeM;
    if (0 <= this->ignoreIndex && this->ignoreIndex < this->sizeN) {
      weightGm.SetValue(this->ignoreIndex, 0);
    }
    // printf("123: %d x,w: %d, %d ,sum: %d \n",tmpTarget, tmpSelect, tmpWeight,
    // sumRes);
    for (int i = 0; i < this->sizeN; i++) {
      for (int j = 0; j < this->sizeM; j++) {
        tmpTarget = targetGm.GetValue(i * sizeM + j);
        tmpWeight = weightGm.GetValue(tmpTarget);
        sumWeight += tmpWeight;
        tmpSelect = xGm.GetValue(i * strideN + (tmpTarget * this->sizeM) + j);
        sumRes += tmpSelect * tmpWeight;
        // printf("target: %d x,w: %d, %d ,sum: %d \n",tmpTarget, tmpSelect,
        // tmpWeight, sumRes);
      }
    }
    if (this->reduction == 2) {  // "mean"
      yGm.SetValue(0, sumRes * -1 / sumWeight);
    } else {
      yGm.SetValue(0, sumRes * -1);
    }
    return;
  }

 private:
  uint32_t reduction;
  int32_t ignoreIndex;
  uint32_t dimNum, sizeN, sizeC, sizeM;

  // AscendC::TPipe pipe;
  AscendC::GlobalTensor<DTYPE_X> xGm;
  AscendC::GlobalTensor<DTYPE_TARGET> targetGm;
  AscendC::GlobalTensor<DTYPE_WEIGHT> weightGm;
  AscendC::GlobalTensor<DTYPE_Y> yGm;
};